import torch
import torch.nn as nn


class MF(nn.Module):
    def __init__(self, hyper_params):
        super(MF, self).__init__()
        self.hyper_parmas = hyper_params

        self.item_embedding = nn.Embedding(
            num_embeddings=hyper_params['item_cnt'],
            embedding_dim=hyper_params['embed_dim']
        )
        self.user_embedding = nn.Embedding(
            num_embeddings=hyper_params['user_cnt'],
            embedding_dim=hyper_params['embed_dim']
        )

        self.apply_extra = True
        if 'apply_extra' in hyper_params:
            self.apply_extra = hyper_params['apply_extra']

        # Set Weight of Extra Info
        if self.apply_extra:
            assert 'extra_weight' in hyper_params and 'extra_dim' in hyper_params

            self.extra_weight = hyper_params['extra_weight']
            self.extra_var = nn.Parameter(
                data=torch.FloatTensor(
                    hyper_params['user_cnt'], hyper_params['extra_dim'], hyper_params['embed_dim']),
                requires_grad=True
            )
            self.register_parameter(name='extra_var', param=self.extra_var)
            nn.init.normal_(self.extra_var)

    def forward(self, item_id, user_id, extra_info=None):
        item_embed = self.item_embedding(item_id)
        user_embed = self.user_embedding(user_id)

        if self.apply_extra:
            assert extra_info is not None and extra_info.shape[1] == self.hyper_parmas['extra_dim']
            extra_embed = (extra_info.unsqueeze(
                1) @ self.extra_var[user_id]).squeeze(1)
            user_embed = user_embed + extra_embed * self.extra_weight

        result = torch.sum(item_embed * user_embed, dim=-1)

        return result

    def get_user_embed(self, user_id, extra_info=None):
        user_embed = self.user_embedding(user_id)

        if self.apply_extra:
            assert extra_info is not None and extra_info.shape[1] == self.hyper_parmas['extra_dim']
            extra_embed = (extra_info.unsqueeze(
                1) @ self.extra_var[user_id]).squeeze(1)
            user_embed = user_embed + extra_embed * self.extra_weight

        return user_embed

    def get_item_embed(self, item_id):
        item_embed = self.item_embedding(item_id)

        return item_embed


class PairwiseLoss:
    def __init__(self, hyper_params, mf):
        self.hyper_params = hyper_params
        self.apply_weight = True

        if 'apply_weight' in hyper_params:
            self.apply_weight = hyper_params['apply_weight']

        self.MF = mf

    def get_loss(self, pos_item, neg_item, item_weight, user_id, pos_extra=None, neg_extra=None):
        pos_val = self.MF(pos_item, user_id, pos_extra)
        neg_val = self.MF(neg_item, user_id, neg_extra)

        result = pos_val - neg_val
        if self.apply_weight:
            result = result * item_weight

        loss = -torch.mean(torch.log(torch.sigmoid(result)+1e-7))

        return loss
